import numpy as np

from jili.calc import calcor_base
import pandas as pd
class std_meanupdown(calcor_base):
    def __init__(self,ta_arg={}):
        self.out = ["std_up","std_down"]
        self.input = ["close"]
        self.parameters = {"timeperiod":20}
        super().__init__(ta_arg)
        self.batch=self.parameters["timeperiod"]
        self.lookback = self.batch - 1
    def calc(self):
        data=self.hisbars[self.input[0]]
        mean=data.mean()
        up=data[data>mean]-mean
        down=data[data<mean]-mean
        up=up**2
        up=up.sum()
        up=(up/self.batch)**(1/2)
        down = down ** 2
        down = down.sum()
        down = (down / self.batch) ** (1 / 2)
        return up,down
class std_0updown(calcor_base):
    def __init__(self,ta_arg={}):
        self.out = ["std_up","std_down"]
        self.input = ["close"]
        self.parameters = {"timeperiod":20}
        super().__init__(ta_arg)
        self.batch=self.parameters["timeperiod"]
    def calc(self):
        data=self.hisbars[self.input[0]]
        if pd.isna(data[0]):
            return None,None
        else:
            up=data[data>0]
            down=data[data<0]
            up = up ** 2
            up = up.sum()
            up = (up / self.batch) ** (1 / 2)
            down = down ** 2
            down = down.sum()
            down = (down / self.batch) ** (1 / 2)
            return up,down
class std_self(calcor_base):
    def __init__(self,ta_arg={}):
        self.out = ["std"]
        self.input = ["close"]
        self.parameters = {"timeperiod":20}
        super().__init__(ta_arg)
        self.batch=self.parameters["timeperiod"]
    def calc(self):
        data=self.hisbars[self.input[0]]
        if pd.isna(data[0]):
            return None
        else:
            std=np.std(data)
            return std